from repo.yolov5.hubconf import yolov5m # 导入完整路径防止在同一个环境中运行多个python function时的路径查找错误

# 载入模型
class yolo:
    # 构造
    def __init__(self):
        pass

    # 载入模型
    def load_model(self):
        self.model = yolov5m()
        self.model = self.model.cuda().eval()

    # [Mat,Mat,...] -> [[Rect,Rect,...],[Rect,Rect,...],...]
    def inference(self, data):
        self.res = self.model(data)
        self.res.print()
        # 将推理结果变成cv::Rect
        self.res = self.res.tolist()
        for i in range(len(self.res)):
            self.res[i] = self.res[i].xywhn.cpu()
        # 筛选出所有的车
        target_classes = [2,5,7] # 2为car,5为bus,7为truck
        rectangles_list = []
        # 对每一张图片
        for res in self.res:
            rectangles = []
            # 对识别出的每一个物体
            for i in res:
                x,y,w,h,n,c = i[0],i[1],i[2],i[3],i[4],(int)(i[5])
                if y >= 0.5 and y <= 0.98 and x >= 0.02 and x <= 0.8 and (c in target_classes): # 取划定范围内的车辆
                    # 封装成cv::Rect
                    rectangle = (x.item()-w.item()/2, y.item()-h.item()/2, w.item(), h.item())
                    rectangles.append(rectangle)
            rectangles_list.append(rectangles)
        return rectangles_list